#!/usr/bin/env python3
"""
ArUco标记对象类

该模块定义了 ArucoMarker 类，用于表示和管理单个ArUco标记的信息。
"""

import numpy as np
from typing import Optional, Tuple, List
import time


class ArucoMarker:
    """
    ArUco标记对象
    
    用于存储和管理单个ArUco标记的所有相关信息，包括：
    - 标记的基本属性（ID、名称、大小）
    - 当前的位置和姿态信息
    - 跟踪状态和历史信息
    """
    
    def __init__(self, name: str, marker_id: int, size_mm: float, window_size: int = 5):
        """
        初始化ArUco标记对象
        
        Args:
            name (str): 用户自定义的标记名称
            marker_id (int): ArUco标记的ID
            size_mm (float): 标记的实际大小，单位为毫米
            window_size (int): 滑动窗口滤波的窗口长度，默认为5
        """
        # 基本属性
        self.name = name
        self.marker_id = marker_id
        self.size_mm = size_mm
        self.window_size = window_size
        
        # 滑动窗口数据
        self.position_window = []    # 位置滑动窗口
        self.rotation_window = []    # 姿态滑动窗口
        self.distance_window = []    # 距离滑动窗口
        self.velocity_window = []    # 速度滑动窗口
        
        # 当前状态
        self.is_tracked = False  # 是否正在被跟踪
        self.position = None     # 当前位置 (x, y, z) 毫米
        self.rotation = None     # 当前姿态 (roll, pitch, yaw) 度
        self.corners = None      # 当前检测到的角点
        self.distance = None     # 距离相机的距离
        
        # 跟踪历史
        self.lost_frames = 0     # 连续丢失的帧数
        self.total_detections = 0  # 总检测次数
        self.last_seen_time = None  # 最后一次检测到的时间
        self.detection_history = []  # 检测历史记录
        
        # 运动状态
        self.velocity = None     # 速度向量
        self.angular_velocity = None  # 角速度
        self.last_position = None
        self.last_rotation = None
        self.last_update_time = None
        
        print(f"创建ArUco标记: {name} (ID: {marker_id}, 大小: {size_mm}mm, 滤波窗口: {window_size})")
    
    def update_detection(self, position: np.ndarray, rotation: np.ndarray, 
                        corners: np.ndarray, distance: float) -> None:
        """
        更新标记的检测信息
        
        Args:
            position (np.ndarray): 位置向量 [x, y, z] 毫米
            rotation (np.ndarray): 旋转角度 [roll, pitch, yaw] 度
            corners (np.ndarray): 检测到的角点坐标
            distance (float): 距离相机的距离
        """
        current_time = time.time()
        
        # 将新的检测数据添加到滑动窗口
        self.position_window.append(position.copy())
        self.rotation_window.append(rotation.copy())
        self.distance_window.append(distance)
        
        # 保持窗口大小
        if len(self.position_window) > self.window_size:
            self.position_window.pop(0)
            self.rotation_window.pop(0)
            self.distance_window.pop(0)
        
        # 计算滤波后的位置、姿态和距离
        filtered_position = np.mean(self.position_window, axis=0)
        filtered_rotation = np.mean(self.rotation_window, axis=0)
        filtered_distance = np.mean(self.distance_window)
        
        # 计算运动信息
        current_velocity = None
        if self.is_tracked and self.last_update_time is not None:
            dt = current_time - self.last_update_time
            if dt > 0:
                # 计算位置速度
                if self.last_position is not None:
                    current_velocity = (filtered_position - self.last_position) / dt
                
                # 计算角速度
                if self.last_rotation is not None:
                    angle_diff = filtered_rotation - self.last_rotation
                    # 处理角度跳跃
                    angle_diff = np.where(angle_diff > 180, angle_diff - 360, angle_diff)
                    angle_diff = np.where(angle_diff < -180, angle_diff + 360, angle_diff)
                    self.angular_velocity = angle_diff / dt
        
        # 将速度添加到滑动窗口并进行滤波
        if current_velocity is not None:
            self.velocity_window.append(current_velocity.copy())
            # 保持窗口大小
            if len(self.velocity_window) > self.window_size:
                self.velocity_window.pop(0)
            # 计算滤波后的速度
            self.velocity = np.mean(self.velocity_window, axis=0)
        
        # 保存当前状态作为历史
        self.last_position = self.position.copy() if self.position is not None else None
        self.last_rotation = self.rotation.copy() if self.rotation is not None else None
        self.last_update_time = current_time
        
        # 更新当前状态（使用滤波后的数据）
        self.position = filtered_position
        self.rotation = filtered_rotation
        self.corners = corners.copy()
        self.distance = filtered_distance
        self.is_tracked = True
        self.lost_frames = 0
        self.total_detections += 1
        self.last_seen_time = current_time
        
        # 添加到历史记录（保持最近100次检测）
        detection_record = {
            'time': current_time,
            'position': position.copy(),
            'rotation': rotation.copy(),
            'distance': distance
        }
        self.detection_history.append(detection_record)
        if len(self.detection_history) > 100:
            self.detection_history.pop(0)
    
    def update_lost(self) -> None:
        """更新标记丢失状态"""
        self.lost_frames += 1
        # 重置运动信息
        self.velocity = None
        self.angular_velocity = None
    
    def mark_lost(self) -> None:
        """标记为丢失状态"""
        self.is_tracked = False
        self.velocity = None
        self.angular_velocity = None
        
        # 清空滑动窗口数据
        self.position_window.clear()
        self.rotation_window.clear()
        self.distance_window.clear()
        self.velocity_window.clear()
        
        print(f"标记 {self.name} 已丢失 (连续 {self.lost_frames} 帧未检测到)")
        print(f"滑动窗口数据已清空，重新开始数据积累")
    
    def get_status(self) -> dict:
        """
        获取标记的完整状态信息
        
        Returns:
            dict: 包含标记所有状态信息的字典
        """
        return {
            'name': self.name,
            'marker_id': self.marker_id,
            'size_mm': self.size_mm,
            'window_size': self.window_size,
            'current_window_length': len(self.position_window),
            'is_tracked': self.is_tracked,
            'position': self.position.tolist() if self.position is not None else None,
            'rotation': self.rotation.tolist() if self.rotation is not None else None,
            'distance': self.distance,
            'lost_frames': self.lost_frames,
            'total_detections': self.total_detections,
            'last_seen_time': self.last_seen_time,
            'velocity': self.velocity.tolist() if self.velocity is not None else None,
            'angular_velocity': self.angular_velocity.tolist() if self.angular_velocity is not None else None,
            'corners': self.corners.tolist() if self.corners is not None else None
        }
    
    def get_position_3d(self) -> Optional[np.ndarray]:
        """
        获取3D位置向量
        
        Returns:
            np.ndarray or None: 位置向量 [x, y, z] 或 None（如果未跟踪）
        """
        if self.is_tracked and self.position is not None:
            return self.position.copy()
        return None
    
    def get_rotation_euler(self) -> Optional[np.ndarray]:
        """
        获取欧拉角度
        
        Returns:
            np.ndarray or None: 欧拉角 [roll, pitch, yaw] 或 None（如果未跟踪）
        """
        if self.is_tracked and self.rotation is not None:
            return self.rotation.copy()
        return None
    
    def get_transform_matrix(self) -> Optional[np.ndarray]:
        """
        获取4x4变换矩阵
        
        Returns:
            np.ndarray or None: 4x4变换矩阵或None（如果未跟踪）
        """
        if not self.is_tracked or self.position is None or self.rotation is None:
            return None
        
        # 转换欧拉角为旋转矩阵
        roll, pitch, yaw = np.radians(self.rotation)
        
        # 计算旋转矩阵 (ZYX顺序)
        cos_r, sin_r = np.cos(roll), np.sin(roll)
        cos_p, sin_p = np.cos(pitch), np.sin(pitch)
        cos_y, sin_y = np.cos(yaw), np.sin(yaw)
        
        R = np.array([
            [cos_y * cos_p, cos_y * sin_p * sin_r - sin_y * cos_r, cos_y * sin_p * cos_r + sin_y * sin_r],
            [sin_y * cos_p, sin_y * sin_p * sin_r + cos_y * cos_r, sin_y * sin_p * cos_r - cos_y * sin_r],
            [-sin_p, cos_p * sin_r, cos_p * cos_r]
        ])
        
        # 构建4x4变换矩阵
        T = np.eye(4)
        T[:3, :3] = R
        T[:3, 3] = self.position
        
        return T
    
    def is_moving(self, position_threshold: float = 1.0, 
                  rotation_threshold: float = 1.0) -> bool:
        """
        判断标记是否在运动
        
        Args:
            position_threshold (float): 位置变化阈值 (mm/s)
            rotation_threshold (float): 角度变化阈值 (degree/s)
            
        Returns:
            bool: 是否在运动
        """
        if not self.is_tracked or self.velocity is None:
            return False
        
        # 检查位置速度
        position_speed = np.linalg.norm(self.velocity)
        if position_speed > position_threshold:
            return True
        
        # 检查角速度
        if self.angular_velocity is not None:
            angular_speed = np.linalg.norm(self.angular_velocity)
            if angular_speed > rotation_threshold:
                return True
        
        return False
    
    def get_detection_rate(self, time_window: float = 10.0) -> float:
        """
        计算指定时间窗口内的检测率
        
        Args:
            time_window (float): 时间窗口，单位秒
            
        Returns:
            float: 检测率 (0.0 - 1.0)
        """
        if not self.detection_history:
            return 0.0
        
        current_time = time.time()
        recent_detections = [
            d for d in self.detection_history 
            if current_time - d['time'] <= time_window
        ]
        
        if not recent_detections:
            return 0.0
        
        # 估算期望的检测次数（假设30fps）
        expected_detections = time_window * 30
        actual_detections = len(recent_detections)
        
        return min(actual_detections / expected_detections, 1.0)
    
    def get_filter_status(self) -> dict:
        """
        获取滤波器状态信息
        
        Returns:
            dict: 包含滤波器状态的字典
        """
        return {
            'window_size': self.window_size,
            'position_window_length': len(self.position_window),
            'rotation_window_length': len(self.rotation_window),
            'distance_window_length': len(self.distance_window),
            'velocity_window_length': len(self.velocity_window),
            'is_filter_ready': len(self.position_window) >= self.window_size
        }
    
    def reset_filter(self) -> None:
        """
        重置滤波器，清空所有滑动窗口数据
        """
        self.position_window.clear()
        self.rotation_window.clear()
        self.distance_window.clear()
        self.velocity_window.clear()
        print(f"标记 {self.name} 的滤波器已重置")
    
    def __str__(self) -> str:
        """字符串表示"""
        status = "跟踪中" if self.is_tracked else f"丢失({self.lost_frames}帧)"
        if self.position is not None:
            pos_str = f"位置:({self.position[0]:.1f}, {self.position[1]:.1f}, {self.position[2]:.1f})"
        else:
            pos_str = "位置:未知"
        
        window_info = f"滤波窗口:{len(self.position_window)}/{self.window_size}"
        
        return f"ArucoMarker('{self.name}', ID:{self.marker_id}, {status}, {pos_str}, {window_info})"
    
    def __repr__(self) -> str:
        """调试表示"""
        return f"ArucoMarker(name='{self.name}', marker_id={self.marker_id}, size_mm={self.size_mm}, is_tracked={self.is_tracked})"
